package plug

import (
	"context"
	"fmt"
	"reflect"
	"strings"
	"sync"

	"github.com/longbridgeapp/sqlparser"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/schema"
)

const (
	NameEncryption    = "gorm:encryption"
	NameEncryptDest   = "gorm:encryptDest"
	NameEncryptWhere  = "gorm:encryptWhere"
	NameEncryptUpdate = "gorm:encryptUpdate"
	NameDecryptDest   = "gorm:decryptDest"

	ClauseSelect = "SELECT"
	ClauseFrom   = "FROM"
	ClauseWhere  = "WHERE"
)

type gormEncryptionContextKey struct{}

var (
	encryptWhereKey = &gormEncryptionContextKey{}
)

// Encryptor contains Encrypt and Decrypt method
type Encryptor interface {
	Encrypt(input string) (*string, error)
	Decrypt(input string) (*string, error)
}

// ColumnDict stores encryption table
type ColumnDict struct {
	// Info: key is tableName, value is map[string]int32, which stores columnNames
	Dict map[string]map[string]int32
}

// Encryption plug encrypts column value before insert or update, and decrypts column value after query
type Encryption struct {
	*gorm.DB
	ConnPool   *ConnPool
	querys     sync.Map
	ColumnDict *ColumnDict
	Encryptor  Encryptor
	_tables    []interface{}
}

// Name plugin name for Gorm plugin interface
func (s *Encryption) Name() string {
	return NameEncryption
}

// Initialize implement for Gorm plugin interface
func (s *Encryption) Initialize(db *gorm.DB) error {
	s.DB = db
	s.registerCallbacks(db)

	return s.compile()
}

func (s *Encryption) compile() error {
	return nil
}

func (s *Encryption) registerCallbacks(db *gorm.DB) {
	// gorm actions contains: create, query, update, delete, row, raw
	s.Callback().Create().Before("*").Register(NameEncryptDest, s.encrypt)
	s.Callback().Query().Before("*").Register(NameEncryptWhere, s.encryptWhere)
	s.Callback().Query().After("*").Register(NameDecryptDest, s.decrypt)
	s.Callback().Update().Before("*").Register(NameEncryptWhere, s.encryptUpdate)
	s.Callback().Delete().Before("*").Register(NameEncryptWhere, s.encryptWhere)
	// These two methods are treated as gorm Query method for now
	s.Callback().Row().Before("*").Register(NameEncryptWhere, s.encryptWhere)
	s.Callback().Raw().Before("*").Register(NameEncryptWhere, s.encryptWhere)
}

func (s *Encryption) switchConn(db *gorm.DB) {
	s.ConnPool = &ConnPool{
		ConnPool:   db.Statement.ConnPool,
		encryption: s,
	}
	db.Statement.ConnPool = s.ConnPool
}

func (s *Encryption) containsTable(tableName string) bool {
	_, ok := s.ColumnDict.Dict[tableName]
	return ok
}

func (s *Encryption) containsColumn(tableName string, columnName string) bool {
	dict, ok := s.ColumnDict.Dict[tableName]
	if !ok {
		return false
	}

	_, ok = dict[columnName]
	return ok
}

func (s *Encryption) encrypt(db *gorm.DB) {
	tableName := db.Statement.Schema.Table
	colMap, ok := s.ColumnDict.Dict[tableName]
	if !ok {
		return
	}

	// Acquire Model to insert or update
	dest := db.Statement.Dest
	// Need to get element twice
	destValue := reflect.ValueOf(dest)
	destValueKind := destValue.Kind()
	var destStruct reflect.Value
	if reflect.Ptr == destValueKind {
		destStruct = destValue.Elem()
	} else {
		destStruct = destValue
	}
	elemKind := destStruct.Kind()

	for k := range colMap {
		colName := k
		field, ok := db.Statement.Schema.FieldsByDBName[colName]
		if !ok {
			continue
		}

		if schema.String != field.DataType {
			continue
		}

		if reflect.Slice == elemKind {
			elem := destStruct
			num := elem.Len()
			for i := 0; i < num; i++ {
				item := elem.Index(i)
				if reflect.Ptr == item.Kind() {
					itemVal := item.Elem()
					s.encryptValue(db, &itemVal, field.Name, colName, tableName)
				} else if reflect.Struct == item.Kind() {
					s.encryptValue(db, &item, field.Name, colName, tableName)
				} else {
					// impossible case
				}
			}
		} else if reflect.Ptr == elemKind {
			elem := destStruct.Elem()
			s.encryptValue(db, &elem, field.Name, colName, tableName)
		} else if reflect.Struct == elemKind {
			s.encryptValue(db, &destStruct, field.Name, colName, tableName)
		} else {
			continue
		}
	}
}

func handleRawValue(db *gorm.DB, s *Encryption, obj *reflect.Value, fieldName, colName, tableName string, handleFunc func(string) (*string, error)) {
	var ok bool
	elem := *obj
	// Get model property by name
	fieldValue := elem.FieldByName(fieldName)
	if !fieldValue.CanSet() {
		msg := fmt.Sprintf("Field value cannot be set, tableName: %s, columnName: %s", tableName, colName)
		s.Logger.Warn(db.Statement.Context, msg)
		return
	}

	var rawVal string
	rawVal, ok = fieldValue.Interface().(string)
	if !ok {
		valPtr, ok := fieldValue.Interface().(*string)
		if !ok {
			msg := fmt.Sprintf("Only string field supports encryption, tableName: %s, columnName: %s", tableName, colName)
			s.Logger.Error(db.Statement.Context, msg)
			return
		}

		if valPtr == nil {
			return
		}
		rawVal = *valPtr
	}

	handledValue, err := handleFunc(rawVal)
	if err != nil {
		msg := fmt.Sprintf("Encryption fails, tableName: %s, columnName: %s, rawVal: %s", tableName, colName, rawVal)
		s.Logger.Error(db.Statement.Context, msg)
		return
	}

	encryptedVal := reflect.ValueOf(handledValue)
	fieldValue.Set(encryptedVal)
	obj = &elem
}

func (s *Encryption) encryptValue(db *gorm.DB, obj *reflect.Value, fieldName, colName, tableName string) {
	handleRawValue(db, s, obj, fieldName, colName, tableName, s.Encryptor.Encrypt)
}

func (s *Encryption) decryptValue(db *gorm.DB, obj *reflect.Value, fieldName, colName, tableName string) {
	handleRawValue(db, s, obj, fieldName, colName, tableName, s.Encryptor.Decrypt)
}

func (s *Encryption) decrypt(db *gorm.DB) {
	tableName := db.Statement.Schema.Table
	colMap, ok := s.ColumnDict.Dict[tableName]
	if !ok {
		return
	}

	// Acquire Model to insert
	dest := db.Statement.Dest
	// Need to get element twice
	destValue := reflect.ValueOf(dest)
	destValueKind := destValue.Kind()
	var destStruct reflect.Value
	if reflect.Ptr == destValueKind {
		destStruct = destValue.Elem()
	} else {
		destStruct = destValue
	}
	elemKind := destStruct.Kind()

	for k := range colMap {
		colName := k
		field, ok := db.Statement.Schema.FieldsByDBName[colName]
		if !ok {
			continue
		}

		if schema.String != field.DataType {
			continue
		}

		if reflect.Slice == elemKind {
			elem := destStruct
			num := elem.Len()
			for i := 0; i < num; i++ {
				item := elem.Index(i)
				if reflect.Ptr == item.Kind() {
					itemVal := item.Elem()
					s.decryptValue(db, &itemVal, field.Name, colName, tableName)
				} else if reflect.Struct == item.Kind() {
					s.decryptValue(db, &item, field.Name, colName, tableName)
				} else {
					// impossible case
				}
			}

		} else if reflect.Struct == elemKind {
			elem := destStruct
			s.decryptValue(db, &elem, field.Name, colName, tableName)
		} else {
			continue
		}
	}
}

func (s *Encryption) encryptWhere(db *gorm.DB) {
	stmt := db.Statement
	finishedVar := stmt.Context.Value(encryptWhereKey)
	finished, ok := finishedVar.(bool)
	if ok && finished {
		return
	}

	where := stmt.Clauses[ClauseWhere]
	exps := where.Expression
	whereClause, ok := exps.(clause.Where)
	if !ok {
		return
	}

	tableName := db.Statement.Schema.Table
	colMap, ok := s.ColumnDict.Dict[tableName]
	if !ok {
		return
	}

	encryptedExprs := make([]clause.Expression, 0, len(whereClause.Exprs))
	for _, expVar := range whereClause.Exprs {
		expr, ok := expVar.(clause.Expr)
		if !ok {
			continue
		}

		for k := range colMap {
			if !strings.Contains(expr.SQL, k) {
				continue
			}
			num := len(expr.Vars)
			encryptedVars := make([]interface{}, 0, num)
			for _, sqlVar := range expr.Vars {
				var strVar string
				strVar, ok := sqlVar.(string)
				if !ok {
					strPtr, ok := sqlVar.(*string)
					if !ok {
						continue
					}
					strVar = *strPtr
				}

				encryptedVar, err := s.Encryptor.Encrypt(strVar)
				if err != nil {
					db.Logger.Error(stmt.Context, "Encryption error: ", err)
					continue
				}
				encryptedVars = append(encryptedVars, encryptedVar)
			}
			expr.Vars = encryptedVars
		}
		encryptedExprs = append(encryptedExprs, expr)
		whereClause.Exprs = encryptedExprs
	}
	where.Expression = whereClause
	stmt.Clauses[ClauseWhere] = where
	ctx := context.WithValue(stmt.Context, encryptWhereKey, true)
	stmt.Context = ctx
	db.Statement = stmt
}

func (s *Encryption) encryptUpdate(db *gorm.DB) {
	s.encryptWhere(db)
	s.encrypt(db)
}

func replaceOrderByTableName(orderBy []*sqlparser.OrderingTerm, oldName, newName string) []*sqlparser.OrderingTerm {
	for i, term := range orderBy {
		if x, ok := term.X.(*sqlparser.QualifiedRef); ok {
			if x.Table.Name == oldName {
				x.Table.Name = newName
				orderBy[i].X = x
			}
		}
	}

	return orderBy
}
